package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"sort"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

// TCPFlow 存储单个TCP流的信息
type TCPFlow struct {
	SrcIP   string
	SrcPort uint16
	DstIP   string
	DstPort uint16

	// 连接状态
	StartTime time.Time
	LastSeen  time.Time
	Status    string // ESTABLISHED, SYN_SENT, etc.

	// 流量统计
	BytesSent       uint64
	BytesReceived   uint64
	PacketsSent     uint64
	PacketsReceived uint64
	AvgPacketSize   uint64
	// TCP指标
	RetransmittedPackets uint64
	DuplicateAcks        uint64
	LostPackets          uint64
	MinRTT               time.Duration
	MaxRTT               time.Duration
	AvgRTT               time.Duration
	RTTSamples           []time.Duration
	LastSentSeq          uint32
	LastAckSeq           uint32
	LastReceivedSeq      uint32
	LastSentTime         time.Time

	// 应用识别相关
	AppName     string  // 应用名称
	AppCategory string  // 应用类别
	Confidence  float64 // 识别可信度(0-1)
}

// FlowKey 用于标识唯一的TCP流
type FlowKey struct {
	SrcIP   string
	SrcPort uint16
	DstIP   string
	DstPort uint16
}

// 全局状态
var (
	tcpFlows      = make(map[FlowKey]*TCPFlow)
	flowMutex     = &sync.RWMutex{}
	iface         string
	snapshotLen   int32         = 1600
	promiscuous   bool          = false
	timeout       time.Duration = pcap.BlockForever
	printInterval time.Duration = 5 * time.Second
)

func init() {
	flag.StringVar(&iface, "iface", "", "网络接口名称 (例如: eth0, en0)")
	flag.DurationVar(&printInterval, "interval", 5*time.Second, "统计信息打印间隔")
	flag.Parse()

	if iface == "" {
		log.Fatal("必须指定网络接口 (-iface)")
	}
}

// 创建FlowKey，确保顺序一致
func getFlowKey(srcIP net.IP, srcPort uint16, dstIP net.IP, dstPort uint16) FlowKey {
	srcStr := fmt.Sprintf("%s:%d", srcIP, srcPort)
	dstStr := fmt.Sprintf("%s:%d", dstIP, dstPort)

	if srcStr < dstStr {
		return FlowKey{
			SrcIP:   srcIP.String(),
			SrcPort: srcPort,
			DstIP:   dstIP.String(),
			DstPort: dstPort,
		}
	}
	return FlowKey{
		SrcIP:   dstIP.String(),
		SrcPort: dstPort,
		DstIP:   srcIP.String(),
		DstPort: srcPort,
	}
}

// 更新TCP流的RTT信息
func updateRTT(flow *TCPFlow, tcp *layers.TCP, packetTime time.Time) {
	if tcp.ACK && tcp.Ack > 0 && flow.LastSentSeq > 0 && tcp.Ack >= flow.LastSentSeq {
		rtt := packetTime.Sub(flow.LastSentTime)
		if rtt > 0 {
			// 限制样本数量，避免内存增长
			if len(flow.RTTSamples) > 100 {
				flow.RTTSamples = flow.RTTSamples[1:] // 移除最旧的样本
			}
			flow.RTTSamples = append(flow.RTTSamples, rtt)

			// 更新RTT统计
			if flow.MinRTT == 0 || rtt < flow.MinRTT {
				flow.MinRTT = rtt
			}
			if rtt > flow.MaxRTT {
				flow.MaxRTT = rtt
			}

			// 计算平均RTT
			total := time.Duration(0)
			for _, t := range flow.RTTSamples {
				total += t
			}
			flow.AvgRTT = total / time.Duration(len(flow.RTTSamples))
		}
	}
}

// 从TLS握手数据中提取SNI（服务器名称指示）
func extractTLSSNI(payload []byte) (string, bool) {
	if len(payload) < 5 || payload[0] != 0x16 { // 不是TLS握手包
		return "", false
	}

	// 跳过记录头(5字节)
	pos := 5
	if pos+1 > len(payload) {
		return "", false
	}

	// 检查握手类型和长度
	handshakeType := payload[pos]
	pos++
	if handshakeType != 0x01 { // 不是客户端问候
		return "", false
	}

	// 解析握手长度
	if pos+3 > len(payload) {
		return "", false
	}
	handshakeLen := int(binary.BigEndian.Uint32(payload[pos:]))
	pos += 3
	if pos+handshakeLen > len(payload) {
		return "", false
	}

	// 跳过版本和随机数
	pos += 2 + 32 // 2字节版本 + 32字节随机数
	if pos+1 > len(payload) {
		return "", false
	}

	// 跳过会话ID
	sessionIDLen := int(payload[pos])
	pos++
	pos += sessionIDLen
	if pos+2 > len(payload) {
		return "", false
	}

	// 跳密码套件
	cipherSuitesLen := int(binary.BigEndian.Uint16(payload[pos:]))
	pos += 2
	pos += cipherSuitesLen
	if pos+1 > len(payload) {
		return "", false
	}

	// 跳过压缩方法
	compressionMethodsLen := int(payload[pos])
	pos++
	pos += compressionMethodsLen
	if pos+2 > len(payload) {
		return "", false
	}

	// 解析扩展
	extensionsLen := int(binary.BigEndian.Uint16(payload[pos:]))
	pos += 2
	end := pos + extensionsLen

	for pos < end {
		if pos+4 > len(payload) {
			break
		}
		extType := binary.BigEndian.Uint16(payload[pos:])
		extLen := int(binary.BigEndian.Uint16(payload[pos+2:]))
		pos += 4

		if extType == 0 { // SNI扩展
			if pos+2 > len(payload) {
				break
			}
			sniListLen := int(binary.BigEndian.Uint16(payload[pos:]))
			pos += 2
			sniEnd := pos + sniListLen

			for pos < sniEnd {
				if pos+3 > len(payload) {
					break
				}
				nameType := payload[pos]
				nameLen := int(binary.BigEndian.Uint16(payload[pos+1:]))
				pos += 3

				if nameType == 0 && pos+nameLen <= len(payload) {
					return string(payload[pos : pos+nameLen]), true
				}
				pos += nameLen
			}
			break
		}
		pos += extLen
	}

	return "", false
}

// 从应用层数据中提取特征
func extractAppFeatures(payload []byte) map[string]string {
	features := make(map[string]string)
	if len(payload) == 0 {
		return features
	}

	// 检测HTTP特征
	firstLine := string(payload[:min(len(payload), 512)])
	if strings.HasPrefix(firstLine, "GET ") || strings.HasPrefix(firstLine, "POST ") ||
		strings.HasPrefix(firstLine, "PUT ") || strings.HasPrefix(firstLine, "DELETE ") ||
		strings.HasPrefix(firstLine, "HEAD ") || strings.HasPrefix(firstLine, "OPTIONS ") ||
		strings.HasPrefix(firstLine, "HTTP/1.") {

		// 提取Host头
		if hostStart := strings.Index(firstLine, "Host: "); hostStart > 0 {
			hostEnd := strings.Index(firstLine[hostStart:], "\r\n")
			if hostEnd > 0 {
				host := strings.TrimSpace(firstLine[hostStart+6 : hostStart+hostEnd])
				features["http_host"] = host
			}
		}

		// 提取User-Agent
		if uaStart := strings.Index(firstLine, "User-Agent: "); uaStart > 0 {
			uaEnd := strings.Index(firstLine[uaStart:], "\r\n")
			if uaEnd > 0 {
				ua := strings.TrimSpace(firstLine[uaStart+12 : uaStart+uaEnd])
				features["http_ua"] = ua
			}
		}
	}

	// 检测TLS特征
	if len(payload) > 5 && payload[0] == 0x16 { // TLS握手标记
		if sni, ok := extractTLSSNI(payload); ok {
			features["tls_sni"] = sni
		}
	}

	return features
}

// 基于端口和协议的基础识别
func identifyByPortProtocol(flow *TCPFlow) (appName string, confidence float64) {
	portAppMap := map[uint16]string{
		80:   "HTTP服务",
		443:  "HTTPS服务",
		21:   "FTP服务",
		22:   "SSH服务",
		53:   "DNS服务",
		3306: "MySQL数据库",
		6379: "Redis服务",
		8080: "HTTP代理/应用服务器",
		25:   "SMTP邮件服务",
		110:  "POP3邮件服务",
		143:  "IMAP邮件服务",
		3389: "远程桌面服务",
		5900: "VNC服务",
	}

	if app, ok := portAppMap[flow.SrcPort]; ok {
		return app, 0.8
	}
	if app, ok := portAppMap[flow.DstPort]; ok {
		return app, 0.8
	}

	return "未知服务", 0.2
}

// 基于域名和特征字符串识别应用
func identifyByDomainFeatures(features map[string]string) (appName, appCategory string, confidence float64) {
	domainMap := map[string][2]string{
		"baidu.com":      {"百度搜索", "网页浏览"},
		"google.com":     {"谷歌搜索", "网页浏览"},
		"weixin.qq.com":  {"微信", "即时通讯"},
		"qq.com":         {"腾讯QQ", "即时通讯"},
		"youtube.com":    {"YouTube", "视频流"},
		"bilibili.com":   {"哔哩哔哩", "视频流"},
		"taobao.com":     {"淘宝", "电子商务"},
		"jd.com":         {"京东", "电子商务"},
		"aliyun.com":     {"阿里云", "云服务"},
		"aws.amazon.com": {"亚马逊云", "云服务"},
		"github.com":     {"GitHub", "代码托管"},
		"gitlab.com":     {"GitLab", "代码托管"},
	}

	var targetDomain string
	if host, ok := features["http_host"]; ok {
		targetDomain = host
	} else if sni, ok := features["tls_sni"]; ok {
		targetDomain = sni
	}

	if targetDomain != "" {
		for domain, appInfo := range domainMap {
			if strings.Contains(targetDomain, domain) {
				return appInfo[0], appInfo[1], 0.9
			}
		}
	}

	if ua, ok := features["http_ua"]; ok {
		switch {
		case strings.Contains(ua, "Chrome"):
			return "Chrome浏览器", "网页浏览", 0.7
		case strings.Contains(ua, "Firefox"):
			return "Firefox浏览器", "网页浏览", 0.7
		case strings.Contains(ua, "Safari"):
			return "Safari浏览器", "网页浏览", 0.7
		case strings.Contains(ua, "Edge"):
			return "Edge浏览器", "网页浏览", 0.7
		case strings.Contains(ua, "Mobile"):
			return "移动设备浏览器", "网页浏览", 0.6
		}
	}

	return "未知应用", "未知类别", 0.3
}

// 基于流量行为特征识别应用类别
func identifyByBehavior(flow *TCPFlow) (appCategory string, confidence float64) {
	duration := flow.LastSeen.Sub(flow.StartTime).Seconds()
	if duration < 1 {
		return "未知类别", 0.2
	}

	totalBytes := flow.BytesSent + flow.BytesReceived
	bytesPerSec := float64(totalBytes) / duration
	packetsPerSec := float64(flow.PacketsSent+flow.PacketsReceived) / duration

	switch {
	case bytesPerSec > 1024*1024 && packetsPerSec > 10:
		return "视频流", 0.7
	case bytesPerSec < 1024*100 && packetsPerSec > 5 && flow.AvgRTT < 500*time.Millisecond:
		return "即时通讯", 0.6
	case flow.AvgPacketSize > 1024*10 && totalBytes > 1024*1024:
		return "文件传输", 0.6
	case bytesPerSec < 1024*500 && packetsPerSec < 100:
		return "网页浏览", 0.5
	}

	return "未知类别", 0.2
}

// 综合识别应用
func identifyApplication(flow *TCPFlow, payload []byte) {
	features := extractAppFeatures(payload)

	// 多维度识别
	portApp, portConf := identifyByPortProtocol(flow)
	domainApp, domainCat, domainConf := identifyByDomainFeatures(features)
	behaviorCat, behaviorConf := identifyByBehavior(flow)

	// 融合结果
	if domainConf >= portConf && domainConf >= behaviorConf {
		flow.AppName = domainApp
		flow.AppCategory = domainCat
		flow.Confidence = domainConf
	} else if portConf >= domainConf && portConf >= behaviorConf {
		flow.AppName = portApp
		flow.AppCategory = "基础服务"
		flow.Confidence = portConf
	} else {
		flow.AppName = "未知应用"
		flow.AppCategory = behaviorCat
		flow.Confidence = behaviorConf
	}
}

// 处理TCP数据包
func handleTCPPacket(packet gopacket.Packet) {
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	if ipLayer == nil {
		return
	}
	ip, _ := ipLayer.(*layers.IPv4)

	tcpLayer := packet.Layer(layers.LayerTypeTCP)
	if tcpLayer == nil {
		return
	}
	tcp, _ := tcpLayer.(*layers.TCP)

	key := getFlowKey(ip.SrcIP, uint16(tcp.SrcPort), ip.DstIP, uint16(tcp.DstPort))
	isOriginalDirection := (key.SrcIP == ip.SrcIP.String() && key.SrcPort == uint16(tcp.SrcPort))

	flowMutex.Lock()
	defer flowMutex.Unlock()

	flow, exists := tcpFlows[key]
	if !exists {
		flow = &TCPFlow{
			SrcIP:     key.SrcIP,
			SrcPort:   key.SrcPort,
			DstIP:     key.DstIP,
			DstPort:   key.DstPort,
			StartTime: packet.Metadata().Timestamp,
			Status:    "NEW",
		}
		tcpFlows[key] = flow
	}

	// 更新最后活动时间
	flow.LastSeen = packet.Metadata().Timestamp

	// 更新TCP状态
	if tcp.SYN && !tcp.ACK {
		flow.Status = "SYN_SENT"
	} else if tcp.SYN && tcp.ACK {
		flow.Status = "SYN_RECEIVED"
	} else if tcp.ACK && flow.Status != "ESTABLISHED" {
		flow.Status = "ESTABLISHED"
	} else if tcp.FIN {
		flow.Status = "FIN_WAIT"
	} else if tcp.RST {
		flow.Status = "RESET"
	}

	// 更新流量统计
	packetSize := uint64(len(packet.Data()))
	if isOriginalDirection {
		flow.BytesSent += packetSize
		flow.PacketsSent++
		flow.LastSentSeq = tcp.Seq
		flow.LastSentTime = packet.Metadata().Timestamp
	} else {
		flow.BytesReceived += packetSize
		flow.PacketsReceived++
		flow.LastReceivedSeq = tcp.Seq
	}

	// 计算平均包大小（总字节数 / 总包数）
	totalBytes := flow.BytesSent + flow.BytesReceived
	totalPackets := flow.PacketsSent + flow.PacketsReceived
	if totalPackets > 0 {
		flow.AvgPacketSize = uint64(float64(totalBytes) / float64(totalPackets))
	}

	if tcp.ACK {
		flow.LastAckSeq = tcp.Ack
	}

	// 检测重传和丢包
	if !tcp.SYN && isOriginalDirection {
		if tcp.Seq <= flow.LastAckSeq && flow.PacketsSent > 1 {
			flow.RetransmittedPackets++
		}
	} else if !tcp.SYN && !isOriginalDirection {
		if tcp.ACK && tcp.Ack == flow.LastAckSeq && flow.PacketsReceived > 1 {
			flow.DuplicateAcks++
			flow.LostPackets++
		}
	}

	// 更新RTT
	updateRTT(flow, tcp, packet.Metadata().Timestamp)

	// 提取应用层数据并识别应用
	applicationLayer := packet.ApplicationLayer()
	if applicationLayer != nil && len(applicationLayer.Payload()) > 0 {
		identifyApplication(flow, applicationLayer.Payload())
	}
}

// 打印统计信息
func printStats() {
	flowMutex.RLock()
	defer flowMutex.RUnlock()

	fmt.Printf("\n===== TCP监控统计 (%s) =====\n", time.Now().Format(time.RFC3339))
	fmt.Printf("总TCP连接数: %d\n", len(tcpFlows))
	fmt.Println("----------------------------------------------------------------------------------------------------------------------------------------------------------------")
	fmt.Printf("%-42s %-12s %-10s %-10s %-8s %-8s %-8s %-10s %-20s %-15s %-6s\n",
		"连接", "状态", "上行(KB)", "下行(KB)", "丢包数", "重传数", "重复ACK", "平均RTT", "应用名称", "应用类别", "可信度")
	fmt.Println("----------------------------------------------------------------------------------------------------------------------------------------------------------------")

	flows := make([]*TCPFlow, 0, len(tcpFlows))
	for _, flow := range tcpFlows {
		flows = append(flows, flow)
	}
	sort.Slice(flows, func(i, j int) bool {
		return flows[i].LastSeen.After(flows[j].LastSeen)
	})

	displayCount := 20
	if len(flows) < displayCount {
		displayCount = len(flows)
	}

	for i := 0; i < displayCount; i++ {
		flow := flows[i]
		connStr := fmt.Sprintf("%s:%d -> %s:%d",
			flow.SrcIP, flow.SrcPort, flow.DstIP, flow.DstPort)
		upKB := float64(flow.BytesSent) / 1024
		downKB := float64(flow.BytesReceived) / 1024

		rttStr := "-"
		if flow.AvgRTT > 0 {
			rttStr = flow.AvgRTT.String()
		}

		fmt.Printf("%-42s %-12s %-10.2f %-10.2f %-8d %-8d %-8d %-10s %-20s %-15s %-6.2f\n",
			connStr,
			flow.Status,
			upKB,
			downKB,
			flow.LostPackets,
			flow.RetransmittedPackets,
			flow.DuplicateAcks,
			rttStr,
			flow.AppName,
			flow.AppCategory,
			flow.Confidence,
		)
	}

	fmt.Println("----------------------------------------------------------------------------------------------------------------------------------------------------------------")
}

// 清理过期连接
func cleanupOldFlows() {
	flowMutex.Lock()
	defer flowMutex.Unlock()

	now := time.Now()
	expirationTime := 5 * time.Minute

	for key, flow := range tcpFlows {
		if now.Sub(flow.LastSeen) > expirationTime {
			delete(tcpFlows, key)
		}
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func main() {
	handle, err := pcap.OpenLive(iface, snapshotLen, promiscuous, timeout)
	if err != nil {
		log.Fatalf("无法打开设备 %s: %v", iface, err)
	}
	defer handle.Close()

	if err := handle.SetBPFFilter("tcp"); err != nil {
		log.Fatalf("无法设置过滤器: %v", err)
	}

	go func() {
		ticker := time.NewTicker(printInterval)
		defer ticker.Stop()

		for range ticker.C {
			printStats()
			cleanupOldFlows()
		}
	}()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packetChan := packetSource.Packets()

	fmt.Printf("开始监控TCP连接 on %s...\n", iface)
	fmt.Printf("每 %v 打印一次统计信息\n", printInterval)
	fmt.Println("按Ctrl+C退出")

	for {
		select {
		case packet := <-packetChan:
			handleTCPPacket(packet)
		case <-sigChan:
			fmt.Println("\n收到退出信号，正在保存最后的统计信息...")
			printStats()
			fmt.Println("监控已停止")
			return
		}
	}
}
